package com.zong.ai.agent.rag;

import org.springframework.ai.document.Document;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * @author: zongzi
 * @description: 游戏王md格式知识库加载器
 * @date: 2025/6/23
 */
@Component
public class YouGiOnMarkDownLoader {
    @Value("${rag.load-path}")
    private String PREFIX;
    private final String SUFFIX = ".md";
    private final ResourcePatternResolver resourcePatternResolver;


    public YouGiOnMarkDownLoader(ResourcePatternResolver resourcePatternResolver) {
        this.resourcePatternResolver = resourcePatternResolver;
    }

    public List<Document> load(String filename) {
        String locationPattern = PREFIX + filename + SUFFIX;
        Resource resource = resourcePatternResolver.getResource(locationPattern);

        MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                .withHorizontalRuleCreateDocument(true)
                .withIncludeCodeBlock(false)
                .withIncludeBlockquote(false)
                .withAdditionalMetadata("filename", filename)
                .build();
        MarkdownDocumentReader reader = new MarkdownDocumentReader(resource, config);
        return reader.read();
    }

    public List<Document> load(Resource resource) {
        String filename = resource.getFilename();

        MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                .withHorizontalRuleCreateDocument(true)
                .withIncludeCodeBlock(false)
                .withIncludeBlockquote(false)
                .withAdditionalMetadata("filename", filename)
                .build();
        MarkdownDocumentReader reader = new MarkdownDocumentReader(resource, config);
        return reader.read();
    }


    public List<Document> loadAll() throws IOException {
        String locationPattern = PREFIX + "*" + SUFFIX;
        Resource[] resources = resourcePatternResolver.getResources(locationPattern);

        List<Document> documents = new ArrayList<>();
        for (Resource resource : resources) {
            List<Document> loads = this.load(resource);
            documents.addAll(loads);
        }
        return documents;
    }
}
